#include "pose_post_process.h"
#include <algorithm>
#include <array>
#include <chrono>
#include <string.h>
#include <vector>
#ifdef LYNXI_PLUGIN
#include "lyn_plugin_dev.h"
#endif

struct PersonPoseConfig {
  std::vector<std::vector<unsigned int>> KPS_COLORS = {
      {0, 255, 0},    {0, 255, 0},    {0, 255, 0},    {0, 255, 0},
      {0, 255, 0},    {255, 128, 0},  {255, 128, 0},  {255, 128, 0},
      {255, 128, 0},  {255, 128, 0},  {255, 128, 0},  {51, 153, 255},
      {51, 153, 255}, {51, 153, 255}, {51, 153, 255}, {51, 153, 255},
      {51, 153, 255}};
  std::vector<std::vector<unsigned int>> SKELETON = {
      {16, 14}, {14, 12}, {17, 15}, {15, 13}, {12, 13}, {6, 12}, {7, 13},
      {6, 7},   {6, 8},   {7, 9},   {8, 10},  {9, 11},  {2, 3},  {1, 2},
      {1, 3},   {2, 4},   {3, 5},   {4, 6},   {5, 7}};
  std::vector<std::vector<unsigned int>> LIMB_COLORS = {
      {51, 153, 255}, {51, 153, 255}, {51, 153, 255}, {51, 153, 255},
      {255, 51, 255}, {255, 51, 255}, {255, 51, 255}, {255, 128, 0},
      {255, 128, 0},  {255, 128, 0},  {255, 128, 0},  {255, 128, 0},
      {0, 255, 0},    {0, 255, 0},    {0, 255, 0},    {0, 255, 0},
      {0, 255, 0},    {0, 255, 0},    {0, 255, 0}};
  std::vector<std::string> class_names = {"person"};
  int class_num = 1;
  int num_point = 17;
  int _anchorLength = 56;
  float _classThreshold = 0.25;
  float _nmsThrehold = 0.45;
  float _kpsVThrehold = 0.5f;
};

struct Rect {
  int x;
  int y;
  int width;
  int height;
};

struct BBox {
  float score;
  Rect rect;
  std::vector<float> kps;
};

int sig[2] = {1, -1};
float result = 5.96046e-08;

float half2float(int16_t ib) {
  int16_t s, e, m;
  s = (ib >> 15) & 0x1;
  e = (ib >> 10) & 0x1f;
  m = ib & 0x3ff;

  // added by puyang.wang@lynxi.com
  {
    if (0 == e)
      return sig[s] * m * result;
    else {
      union {
        unsigned int u32;
        float f32;
      } ou;

      e = (0x1f == e) ? 0xff : (e - 15 + 127);
      ou.u32 = (s << 31) | (e << 23) | (m << 13);
      return ou.f32;
    }
  }
}

float calculateIOU(const Rect &a, const Rect &b) {
  int x1 = std::max(a.x, b.x);
  int y1 = std::max(a.y, b.y);
  int x2 = std::min(a.x + a.width, b.x + b.width);
  int y2 = std::min(a.y + a.height, b.y + b.height);

  if (x2 < x1 || y2 < y1)
    return 0.0f;

  float intersection = (x2 - x1) * (y2 - y1);
  float areaA = a.width * a.height;
  float areaB = b.width * b.height;

  return intersection / (areaA + areaB - intersection);
}

std::vector<int> nmsBoxes(const std::vector<Rect> &boxes,
                          const std::vector<float> &scores,
                          float scoreThreshold, float nmsThreshold) {
  std::vector<int> indices;
  for (size_t i = 0; i < scores.size(); ++i) {
    if (scores[i] >= scoreThreshold) {
      indices.push_back(i);
    }
  }

  std::sort(indices.begin(), indices.end(),
            [&scores](int a, int b) { return scores[a] > scores[b]; });

  std::vector<int> selected;
  std::vector<bool> active(indices.size(), true);

  for (size_t i = 0; i < indices.size(); ++i) {
    if (!active[i])
      continue;

    selected.push_back(indices[i]);
    for (size_t j = i + 1; j < indices.size(); ++j) {
      if (active[j]) {
        float iou = calculateIOU(boxes[indices[i]], boxes[indices[j]]);
        if (iou >= nmsThreshold) {
          active[j] = false;
        }
      }
    }
  }
  return selected;
}

int lynPosePostProcess(PosePostInfo_t *post_info) {
  int16_t *fp16_data_ptr = nullptr;
  lynPoseInfo *poseInfo = nullptr;
#ifdef LYNXI_PLUGIN
  fp16_data_ptr = (int16_t *)lynPluginGetVirtAddr(post_info->output_tensor);
  if (fp16_data_ptr == nullptr) {
    LOG_PLUGIN_E("get output tensor addr error\n");
    return -1;
  }

  poseInfo = (lynPoseInfo *)lynPluginGetVirtAddr(post_info->poseInfo);
  if (poseInfo == nullptr) {
    LOG_PLUGIN_E("get poseInfo addr error\n");
    return -1;
  }
#else
  poseInfo = post_info->poseInfo;
  fp16_data_ptr = (int16_t *)post_info->output_tensor;
#endif

  const float ratio = post_info->ratio;
  const float padx = post_info->padx;
  const float pady = post_info->pady;
  const int img_width = post_info->img_width;
  const int img_height = post_info->img_height;

  PersonPoseConfig default_pose_config;
  PersonPoseConfig *config = &default_pose_config;

  const int model_height = post_info->model_height; // 56
  const int model_width = post_info->model_width;   // 6300

  std::vector<float> confidences;
  std::vector<Rect> boxes;
  std::vector<std::vector<float>> kpss;

  for (int r = 0; r < model_width; ++r) {
    // 直接获取置信度并转换
    int16_t score_fp16 = fp16_data_ptr[4 * model_width + r];
    float score = half2float(score_fp16);

    if (score > config->_classThreshold) {
      // 转换坐标相关参数
      float x = half2float(fp16_data_ptr[0 * model_width + r]);
      float y = half2float(fp16_data_ptr[1 * model_width + r]);
      float w = half2float(fp16_data_ptr[2 * model_width + r]);
      float h = half2float(fp16_data_ptr[3 * model_width + r]);

      x = (x - padx) / ratio;
      y = (y - pady) / ratio;
      w /= ratio;
      h /= ratio;

      int left = static_cast<int>(std::max(x - w * 0.5f, 0.0f));
      int top = static_cast<int>(std::max(y - h * 0.5f, 0.0f));
      int right = static_cast<int>(std::min(left + w, img_width - 1.0f));
      int bottom = static_cast<int>(std::min(top + h, img_height - 1.0f));

      // 处理关键点
      std::vector<float> kps;
      for (int k = 0; k < config->num_point; ++k) {
        float kx = half2float(fp16_data_ptr[(5 + 3 * k) * model_width + r]);
        float ky = half2float(fp16_data_ptr[(5 + 3 * k + 1) * model_width + r]);
        float ks = half2float(fp16_data_ptr[(5 + 3 * k + 2) * model_width + r]);

        kps.push_back((kx - padx) / ratio); // x
        kps.push_back((ky - pady) / ratio); // y
        kps.push_back(ks);                  // score
      }

      confidences.push_back(score);
      boxes.push_back({left, top, right - left, bottom - top});
      kpss.push_back(kps);
    }
  }

  // 执行NMS（与原代码相同）
  std::vector<int> nms_result = nmsBoxes(
      boxes, confidences, config->_classThreshold, config->_nmsThrehold);

  // 填充结果（与原代码相同）
  poseInfo->boxesNum = nms_result.size();
  for (int i = 0; i < poseInfo->boxesNum; ++i) {
    int idx = nms_result[i];

    poseInfo->boxes[i].xmax = boxes[idx].x + boxes[idx].width;
    poseInfo->boxes[i].xmin = boxes[idx].x;
    poseInfo->boxes[i].ymax = boxes[idx].y + boxes[idx].height;
    poseInfo->boxes[i].ymin = boxes[idx].y;
    poseInfo->boxes[i].score = confidences[idx];
    poseInfo->boxes[i].id = 0;

    strncpy(poseInfo->boxes[i].label, "person",
            sizeof(poseInfo->boxes[i].label) - 1);

    for (size_t p = 0; p < kpss[idx].size(); ++p) {
      poseInfo->points[i].kps[p] = kpss[idx][p];
    }
  }

  return 0;
}